library(patchwork)
library(cregg)
library(dplyr)
library(ggplot2)
library(tikzDevice)

estimate_marginal_means <- function(experiment_data){
  # estimates marginal means for second experiment 
  
  formula_program <- num_prog ~ local_status + party
  formula_pork <- num_pork ~ local_status + party
  
  
  programmatic_outcome <- cregg::cj(
    formula = formula_program,
    data = experiment_data,
    id = ~respondent,
    estimate = "mm"
  )
  
  
  pork_outcome <- cregg::cj(
    formula = formula_pork,
    data = experiment_data,
    id = ~respondent,
    estimate = "mm"
  )
  
  second_experiment_result <- rbind(programmatic_outcome, pork_outcome)
  
  return(second_experiment_result)
  
}

format_cregg_output_for_plot <- function(second_experiment_result){
  
  
  second_experiment_result$outcome <- as.factor(second_experiment_result$outcome)
  levels(second_experiment_result$outcome) <- c("Number of Particularistic Policies Selected", "Number of Programmatic Policies Selected")
  levels(second_experiment_result$level) <- c("Elite", "Local", "Non-Local", "DPP", "LDP", "JCP", "CDPJ")
  
  return(second_experiment_result)
}

make_policy_type_selection_plot <- function(plot_data, main_paper_version = TRUE){
  
  if (main_paper_version){
    plot_data <- plot_data %>% 
      filter(feature == "local_status")
  }
  
  
  policy_selection_plot <- ggplot(data = plot_data, aes(x = estimate, y = level, xmin = lower, xmax = upper)) +
    geom_point() + 
    geom_errorbar(width = 0) + 
    geom_vline(xintercept = 1.5, linetype = "dashed") + 
    facet_wrap(. ~ outcome, ncol = 1) + 
    scale_y_discrete(name = "Type of Candidate") + 
    theme_bw()
  
  return(policy_selection_plot)
  
}


path_to_data <- "data"
formatted_survey_file <- "follow_up_survey_formatted.rds"

path_to_figures <- "figures"

path_to_appendix_figures <- "appendix_plots"

carry_over_assumption_plot_file <- "no_carryover_experiment2.eps"
carry_over_assumption_plot_first_task_file <- "no_carryover_experiment2_first_task.eps"

formatted_data <- readRDS(file.path(path_to_data, formatted_survey_file))

formatted_data_task_split <- formatted_data %>% 
  group_split(task)


estimated_mm_by_task <- lapply(formatted_data_task_split, estimate_marginal_means)
estimated_mm_by_task <- lapply(1:3, function(x) mutate(estimated_mm_by_task[[x]], task_number = x))
estimated_mm_by_task <- lapply(estimated_mm_by_task, format_cregg_output_for_plot)
estimated_mm_by_task <- data.table::rbindlist(estimated_mm_by_task)

first_task_plot <- estimated_mm_by_task %>% filter(task_number == 1) %>% 
  make_policy_type_selection_plot() + 
  scale_x_continuous(name = "Estimate Task 1", limits = c(0.75, 2.))

second_task_plot <- estimated_mm_by_task %>% 
  filter(task_number == 2) %>% 
  make_policy_type_selection_plot() + 
  theme(axis.title.y = element_blank(),
        axis.text.y = element_blank(),
        axis.ticks.y = element_blank(),
        legend.position = "none") + 
  scale_x_continuous(name = "Estimate Task 2", limits = c(0.75, 2.))

third_task_plot <- estimated_mm_by_task %>% 
  filter(task_number == 3) %>% 
  make_policy_type_selection_plot() + 
  theme(axis.title.y = element_blank(),
        axis.text.y = element_blank(),
        axis.ticks.y = element_blank(),
        legend.position = "none") + 
  scale_x_continuous(name = "Estimate Task 3", limits = c(0.75, 2.))

# combines plot with patchwork package 
task_separate_plot_experiment_2 <- first_task_plot + second_task_plot + third_task_plot

ggsave(
  filename = file.path(path_to_figures, path_to_appendix_figures, carry_over_assumption_plot_file),
  plot = task_separate_plot_experiment_2,
  height = 10, 
  width = 8,
  units = "in",
  device = cairo_ps,
  fallback_resolution = 200)

ggsave(
  filename = file.path(path_to_figures, path_to_appendix_figures, carry_over_assumption_plot_first_task_file),
  plot = first_task_plot,
  height = 10, 
  width = 8,
  units = "in",
  device = cairo_ps,
  fallback_resolution = 200)






